%{
##########################################################################
##########################################################################
Peter Maxted

Main file to solve 2-income-state model with IG agents

***This code is based heavily on Achdou et al. (2017) codes***
   Available on Ben Moll's website
   See in particular: "HJB equation, implicit method"
                      "KFE Equation"

##########################################################################
##########################################################################
%}

clearvars
addpath('subfunctions')
addpath('output')


%%-------------------------------------------------------------------------
%%-------------------------------------------------------------------------
%Define Basic Model Setup

    presentBias = 0;
    saveOutput = 1;
    plotFigs = 1;
    
    %- - - - - - - - - - - - - - - -
    %Income Process Calibration
    % Guerrieri and Lorenzoni Calibration (+ Floden and Linde)
    Delta = 1/4;
    rate_y_OU = -log(0.967)/Delta;
    sig_y_OU = sqrt((0.017*2*rate_y_OU)/(1-exp(-2*rate_y_OU*Delta)));

    %Construct transition matrix (imposing reflecting barrier)
    dz = 2*sig_y_OU;
    lambda_L = (rate_y_OU*(dz/2))/dz + (sig_y_OU^2)/(2*dz^2);
    lambda_H = lambda_L;
    y_L = exp(-sig_y_OU);
    y_H = exp(sig_y_OU);
    %Adjust so mean-income = 1
        meanincomepre = (y_L + y_H)/2;
        y_L = y_L / meanincomepre;
        y_H = y_H / meanincomepre;
    
    incomeStruct = struct('y_L', y_L, ...
                          'y_H', y_H, ...
                          'lambda_L', lambda_L, ...
                          'lambda_H', lambda_H);

    %Preferences
    switch presentBias
        % -----------------------------------------------------------------
        % GE cases -- iterate on r
        case 0 %Exponential
            rho = 0.035;
            beta = 1;
            savename = 'beta1';
        case 1 %Sophisticated
            rho = 0.025;
            beta = 0.75;
            savename = 'beta0pt75';
        case 2 %Fully Naive
            rho = 0.0245;
            beta = 0.75;
            betaE = 1;
            savename = 'beta0pt75_betaE1';
    end
    gamma = 2;

    %Other
    borrowingLimit = -1/3;
    assetSupply = 3;
    %- - - - - - - - - - - - - - - -
%%-------------------------------------------------------------------------
%%-------------------------------------------------------------------------


%%-------------------------------------------------------------------------
%%-------------------------------------------------------------------------
%Continuous Time Simulation Calibration

	%- - - - - - - - - - - - - - - -
	%Continuous Time Model Parameters
	assert(y_H >= y_L);
    if ~exist('betaE', 'var')
        betaE = beta; %sophistication
    end
    
    psi = (gamma - (1-betaE))/gamma;
        assert(psi > 0 & psi <=1);

	if gamma == 1
         util = @(x) log(x);
         util_hat = @(x) util(x) - log(betaE) + ((betaE - 1)/betaE);
    else
         util = @(x) (x.^(1-gamma) - 1)/(1-gamma);
         util_hat = @(x) (psi/betaE)*util((1/psi)*x) + ((psi - 1)/betaE);
    end
    prefStruct = struct('rho', rho, ...
                        'beta', beta, ...
                        'betaE', betaE, ...
                        'gamma', gamma, ...
                        'psi', psi, ...
                        'util', util, ...
                        'util_hat', util_hat);
    %- - - - - - - - - - - - - - - -

	%- - - - - - - - - - - - - - - -
	%Simulation and Grid Parameters
	I = 30000;
	xmax = 400;
    Xstruct = buildX(1, borrowingLimit, xmax, I);

    maxIterations = 500;
	crit = 10^(-6);
	Delta = 1000;
    simStruct = struct('maxIterations', maxIterations, ...
                       'crit', crit, ...
                       'Delta', Delta);
    %- - - - - - - - - - - - - - - -
%%-------------------------------------------------------------------------
%%-------------------------------------------------------------------------


%%-------------------------------------------------------------------------
%%-------------------------------------------------------------------------
%Find the Equilibrium
    %Step 1. Solve the household problem for a given r
    %Step 2. Find the steady state distribution using the KF equation
    %Step 3. Check if equilibrium asset demand equals B
    %Step 4. Iterate on r until convergence
    critClearing = 10^(-4);

    rInit = rho/2;
    rMax = rho/beta + 0.01;
    rMin = -0.5;
    assetDemand = 100;
    iter = 1;

    % -----------------------------------------
    while abs(assetDemand - assetSupply) > critClearing
        if assetDemand - assetSupply > 0 && iter > 1
			%Too much saving (r too high)
            rMax = r;
            r = 0.5 * (r + rMin)
        elseif assetDemand - assetSupply < 0 && iter > 1
			%Too little saving (r too low)
            rMin = r;
            r = 0.5 * (r + rMax)
        end

        if iter == 1
            r = rInit
        end

        %Solve HJB and KF equations
        [V, C_, S_, A] = solveHJB(r, 0, incomeStruct, prefStruct, Xstruct, simStruct);

        [assetDemand, g] = solveKF(r, A, incomeStruct, prefStruct, Xstruct, simStruct);
        
        iter = iter + 1;
    end
    % -----------------------------------------
    
    %Check that xmax large enough (roughly)
    consRateLimit = (rho - (1-gamma)*r)/(gamma-(1-beta));
    ybar = (lambda_H*y_L + lambda_L*y_H)/(lambda_L + lambda_H);
    if xmax < ((consRateLimit*(ybar/r) - y_H)/(r-consRateLimit))
        warning('xmax should be set higher')
    end
%%-------------------------------------------------------------------------
%%-------------------------------------------------------------------------


%%-------------------------------------------------------------------------
%%-------------------------------------------------------------------------
%Calculate MPCs and Distributional Properties
Gamma_ = MPC_2Huggett(C_, S_, 1/4, 300, incomeStruct, Xstruct);
MPC_ = (Gamma_(2:end, :, 1) - Gamma_(1:end-1, :, 1)) ./ Xstruct.xjump;
Expected_MPC = sum((g(1:end-1, :) *Xstruct.xjump) .* MPC_) ./ sum((g(1:end-1, :) *Xstruct.xjump));

[wealthDistStats] = distStats_2Huggett(g, S_, Xstruct);
%%-------------------------------------------------------------------------
%%-------------------------------------------------------------------------


%%-------------------------------------------------------------------------
%%-------------------------------------------------------------------------
%Save workspace and plot output
if saveOutput == 1
    save(['output/', savename, '_output.mat'])
end

if plotFigs == 1 && isfile('output/beta1_output.mat') && isfile('output/beta0pt75_output.mat') && isfile('output/beta0pt75_betaE1_output.mat')
    plotOutput_2Huggett
end

